package xyz.labmem.lipx.client.core

import io.netty.channel.ChannelInitializer
import io.netty.channel.socket.SocketChannel
import io.netty.handler.codec.LengthFieldBasedFrameDecoder
import io.netty.handler.timeout.IdleStateHandler
import xyz.labmem.lipx.client.core.handler.LabNettyClientHandler
import xyz.labmem.lipx.core.AppContext
import xyz.labmem.lipx.core.AppContext.Client.rowPool
import xyz.labmem.lipx.core.Status
import xyz.labmem.lipx.core.Status.FAILED
import xyz.labmem.lipx.core.Type.FORWARD
import xyz.labmem.lipx.core.codec.LabMessageDecoder
import xyz.labmem.lipx.core.codec.LabMessageEncoder
import xyz.labmem.lipx.core.entity.ClientConfig
import xyz.labmem.lipx.core.net.TcpConnection
import xyz.labmem.lipx.core.tool.LogInfo

/**
 * @description: main
 * @author: ltty
 * @date: 2022/5/27 13:11
 */
class LabSSHPenetrationClient(
    private var config: ClientConfig
) {
    private lateinit var tcp: TcpConnection

    private var remake = false

    private var remaking = false

    private var remakeCount = 0

    init {
        AppContext.Client.allList[config.id]?.status = if (AppContext.Client.allList[config.id]?.status == FAILED) {
            Status.RE_CONNECT
        } else Status.CONNECTING
    }

    fun connect(call: (() -> Unit)? = null): LabSSHPenetrationClient {
        try {
            tcp = TcpConnection()
            val future = tcp.connect(
                config.serverHost!!,
                config.serverPort!!,
                object : ChannelInitializer<SocketChannel>() {
                    @Throws(Exception::class)
                    override fun initChannel(ch: SocketChannel) {
                        ch.pipeline().addLast(
                            LengthFieldBasedFrameDecoder(Int.MAX_VALUE, 0, 4, 0, 4),
                            LabMessageDecoder(), LabMessageEncoder(),
                            IdleStateHandler(60, 30, 0),
                            LabNettyClientHandler(config, config.type == FORWARD) { b, r, rr, msg ->
                                if (!remake)
                                    remake = b
                                if (r) {
                                    remake = false
                                }
                                if (b) {
                                    remakeCount = 0
                                    AppContext.Client.allList[config.id]?.status = Status.CONNECTED
                                    AppContext.Client.connectList[config.id!!] = this@LabSSHPenetrationClient
                                    remaking = false
                                } else {
                                    close()
                                    LogInfo.appendLogError(msg, config)
                                    release(FAILED)
                                }
                                call?.let { it() }
                            }
                        )
                    }
                })
            // 监听
            future.addListener { if (remake && !remaking) close() }
        } catch (e: Exception) {
            release(FAILED)
            LogInfo.appendLogError("lipx连接错误:${e.message}", config)
            call?.let { it() }
        }
        return this
    }

    fun close() {
        LogInfo.appendLog("正在断开当前连接。。", config)
        try {
            remake = false
            tcp.channel?.close()
            LogInfo.appendLog("连接已断开", config)
        } catch (e: Exception) {
            LogInfo.appendLogError("连接断开错误:${e.message}", config)
        } finally {
            rowPool.execute {
                AppContext.Client.connectList.remove(config.id)
            }
            release(Status.IDLE)
        }
    }

    private fun release(status: Status) {
        rowPool.execute {
            AppContext.Client.allList[config.id]?.status = status
        }
    }

    fun restart(newConfig: ClientConfig) {
        config = newConfig
        AppContext.Client.allList[config.id]?.status = Status.RE_CONNECT
        LogInfo.appendLog("正在断开当前连接。。", config)
        try {
            remake = false
            tcp.channel?.close()
            LogInfo.appendLog("连接已断开，准备重启", config)
            connect()
        } catch (e: Exception) {
            AppContext.Client.allList[config.id]?.status = Status.IDLE
            LogInfo.appendLogError("连接重启错误:${e.message}", config)
        }
    }

}